package solutions.leetcode.hard;

import java.util.Random;

/**
 * @author lizhidong
 * <a href="https://leetcode-cn.com/problems/random-pick-with-weight/">528. 按权重随机选择</a>
 */
public class Solution528 {

    /**
     * 给定一个正整数数组w ，其中w[i]代表下标 i的权重（下标从 0 开始），请写一个函数pickIndex，它可以随机地获取下标 i，选取下标 i 的概率与 w[i] 成正比。
     *
     * 例如，对于 w = [1, 3]，挑选下标 0 的概率为 1 / (1 + 3) = 0.25 （即，25%），而选取下标 1 的概率为 3 / (1 + 3) = 0.75（即，75%）。
     *
     * 也就是说，选取下标 i 的概率为 w[i] / sum(w) 。
     *
     * 输入：
     * ["Solution","pickIndex"]
     * [[[1]],[]]
     * 输出：
     * [null,0]
     * 解释：
     * Solution solution = new Solution([1]);
     * solution.pickIndex(); // 返回 0，因为数组中只有一个元素，所以唯一的选择是返回下标 0。
     */

    static class Solution {
        Random random;
        int[] val;
        long len;

        public Solution(int[] w) {
            random = new Random();
            val = w;
            for (int v : w) {
                len += v;
            }
        }

        public int pickIndex() {
            long step = (long) (random.nextDouble() * len) + 1;
            int p = 0;
            while (step > 0 && p < val.length) {
                if (step <= val[p]) {
                    break;
                } else {
                    step -= val[p++];
                }
            }
            return p;
        }
    }

    public static void main(String[] args) {
        Solution solution = new Solution(new int[] {1, 2, 4,8});
        int[] rst = new int[4];
        for (int i = 0; i < 100000; i++) {
            rst[solution.pickIndex()]++;
        }
        for (int v : rst) {
            System.out.println(v);
        }
    }
}
